Draft: GPT2 training on MCQ med data#1111
Open
mina5rovic wants to merge 40 commits into
Open
Conversation
mina5rovic
changed the title
mina5rovic
force-pushed
the
gpt2-training
branch
from
fd4b676 to
3014cf5
Compare
…k for model to contxt 512
JulienVig
reviewed
Jun 1, 2026
JulienVig
left a comment
JulienVig
left a comment
Collaborator
There was a problem hiding this comment.
The logic of the benchmark and the memorization looks good to me! Just had some performance comments
|
|
||
| const logits = tfModel.predict(inputTensor) as tf.Tensor; | ||
| const logProbs = tf.logSoftmax(logits, -1); | ||
| const arr = await logProbs.array() as number[][][]; |
Collaborator
There was a problem hiding this comment.
You're computing the softmax for every position but you only need the last one and line 74 materializes the whole array while you only need 4. You could rewrite this logic such that you only work with the last position, for example:
const optionLogProbs = tf.tidy(() => {
const logits = tfModel.predict(inputTensor) as tf.Tensor3D; // [1, seqLen, vocab]
const lastLogits = logits
.slice([0, promptTokens.length - 1, 0], [1, 1, -1]) // final position only
.reshape([-1]); // [vocab]
const logProbs = tf.logSoftmax(lastLogits); // [vocab]
return tf.gather(logProbs, continuationTokenIDs); // continuationTokenIDs is an array of the 4 continuations' tokenID
});
const scores = await optionLogProbs.array(); // just 4 values
| return output as tf.Tensor; | ||
| }); | ||
|
|
||
| console.log("logits shape:", logits.shape); |
Collaborator
There was a problem hiding this comment.
Make sure to remove the debug prints before merging the PR
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.